
''''
计算原始数据的平均云顶高度并输出NC
'''
import numpy as np
import netCDF4 as nc
import os
import glob
from tqdm import tqdm

# 配置参数
# INPUT_DIR = "/mnt/datastore/liudddata/fy_4Adata/FY20200101test/fy_cth_20200101"
# INPUT_DIR = "/mnt/datastore/liudddata/FY20200101test/fy_cth_202003_01_10"
INPUT_DIR = "/mnt/datastore/liudddata/FY20200101test/result_202003_01_10"
OUTPUT_NC = "/mnt/datastore/liudddata/season/average_fycth_ocean_202003_01_10.nc"
# VARIABLE_NAME = "CTH"
VARIABLE_NAME = "cth"
FILL_VALUE = np.nan

def process_2d_grid():
    # 获取文件列表
    # file_list = sorted(glob.glob(os.path.join(INPUT_DIR, "*_4000M_V0001.NC")))
    file_list = sorted(glob.glob(os.path.join(INPUT_DIR, "*predicted_2d.nc")))
    if not file_list:
        raise ValueError("未找到输入文件")

    # 读取第一个文件获取网格信息
    with nc.Dataset(file_list[0], 'r') as template_ds:
        # 验证维度结构
        if 'y' not in template_ds.dimensions or 'x' not in template_ds.dimensions:
            raise RuntimeError("原始文件维度不符合预期结构")
        # 获取维度大小
        y_dim = len(template_ds.dimensions['y'])
        x_dim = len(template_ds.dimensions['x'])

    # 加载地理坐标数据
    coord_file_name = '/home/liudd/data_preprocessing/FY4A_coordinates.nc'
    with nc.Dataset(coord_file_name, 'r') as coord_ds:
        lat = coord_ds.variables['lat'][:, :].T  # 转置确保维度匹配
        lon = coord_ds.variables['lon'][:, :].T
        # 检查坐标维度与数据维度是否匹配
        if lat.shape != (y_dim, x_dim) or lon.shape != (y_dim, x_dim):
            raise RuntimeError("坐标维度与数据维度不匹配")

    # 初始化累加数组
    sum_data = np.zeros((y_dim, x_dim), dtype=np.float64)
    count_data = np.zeros((y_dim, x_dim), dtype=np.uint32)

    # 逐文件处理
    for file_path in tqdm(file_list, desc="处理进度"):
        with nc.Dataset(file_path, 'r') as ds:
            # 验证数据维度
            if ds[VARIABLE_NAME].shape != (y_dim, x_dim):
                continue
            data = ds[VARIABLE_NAME][:, :].astype(np.float32)
            valid_mask = (~np.isclose(data, FILL_VALUE)) & (~np.isnan(data)) & (data > 0)
            sum_data += np.where(valid_mask, data, 0)
            count_data += valid_mask.astype(np.uint32)

    # 计算平均值
    with np.errstate(divide='ignore', invalid='ignore'):
        mean_data = np.divide(sum_data, count_data, where=count_data > 0)
    mean_data = np.where(count_data > 0, mean_data, FILL_VALUE).astype(np.float32)

    # 创建输出文件
    with nc.Dataset(OUTPUT_NC, 'w', format='NETCDF4') as out_ds:
        # 定义维度
        out_ds.createDimension('y', y_dim)
        out_ds.createDimension('x', x_dim)

        # 写入经纬度变量
        lat_var = out_ds.createVariable('lat', 'f4', ('y', 'x'), zlib=True)
        lon_var = out_ds.createVariable('lon', 'f4', ('y', 'x'), zlib=True)
        lat_var[:, :] = lat
        lon_var[:, :] = lon
        lat_var.units = "degrees_north"
        lat_var.long_name = "latitude"
        lon_var.units = "degrees_east"
        lon_var.long_name = "longitude"

        # 写入平均云顶高度
        cth_var = out_ds.createVariable(
            'mean_cth', 'f4', ('y', 'x'),
            fill_value=FILL_VALUE,
            zlib=True,
            complevel=2
        )
        cth_var[:, :] = mean_data

        # 写入有效数据计数
        valid_count_var = out_ds.createVariable(
            'valid_count', 'u4', ('y', 'x'),
            zlib=True,
            complevel=2
        )
        valid_count_var[:, :] = count_data

        # 添加全局属性
        out_ds.title = "FY-4A网格年平均云顶高度"
        out_ds.grid_type = "不规则二维网格"
        out_ds.source_files = ";".join([os.path.basename(f) for f in file_list])

if __name__ == "__main__":
    try:
        process_2d_grid()
        print(f"处理完成，结果保存至：{OUTPUT_NC}")
    except Exception as e:
        print(f"处理失败：{str(e)}")